/* Implementation derived from:
 * "Quickly Generating Billion-Record Synthetic Databases", Jim Gray et al,
 * SIGMOD 1994
 *
 * The zipfian_int_distribution class is intended to be compatible with other
 * distributions introduced in #include <random> by the C++11 standard.
 *
 * Usage example:
 * #include <random>
 * #include "zipfian_int_distribution.h"
 * int main()
 * {
 *   std::default_random_engine generator;
 *   zipfian_int_distribution<int> distribution(1, 10, 0.99);
 *   int i = distribution(generator);
 * }
 */

/*
 * IMPORTANT: constructing the distribution object requires calculating the zeta
 * value which becomes prohibetively expensive for very large ranges. As an
 * alternative for such cases, the user can pass the pre-calculated values and
 * avoid the calculation every time.
 *
 * Usage example:
 * #include <random>
 * #include "zipfian_int_distribution.h"
 * int main()
 * {
 *   std::default_random_engine generator;
 *   zipfian_int_distribution<int>::param_type p(1, 1e6, 0.99, 27.000);
 *   zipfian_int_distribution<int> distribution(p);
 *   int i = distribution(generator);
 * }
 */

#include <cassert>
#include <cmath>
#include <limits>
#include <random>

template <typename _IntType = int>
class zipfian_int_distribution
{
    static_assert(std::is_integral<_IntType>::value, "Template argument not an integral type.");

public:
    /** The type of the range of the distribution. */
    typedef _IntType result_type;
    /** Parameter type. */
    struct param_type
    {
        typedef zipfian_int_distribution<_IntType> distribution_type;

        explicit param_type(_IntType __a = 0, _IntType __b = std::numeric_limits<_IntType>::max(), double __theta = 0.99)
            : _M_a(__a), _M_b(__b), _M_theta(__theta),
              _M_zeta(zeta(_M_b - _M_a + 1, __theta)), _M_zeta2theta(zeta(2, __theta))
        {
            assert(_M_a <= _M_b && _M_theta > 0.0 && _M_theta < 1.0);
        }

        explicit param_type(_IntType __a, _IntType __b, double __theta, double __zeta)
            : _M_a(__a), _M_b(__b), _M_theta(__theta), _M_zeta(__zeta),
              _M_zeta2theta(zeta(2, __theta))
        {
            __glibcxx_assert(_M_a <= _M_b && _M_theta > 0.0 && _M_theta < 1.0);
        }

        result_type a() const { return _M_a; }

        result_type b() const { return _M_b; }

        double theta() const { return _M_theta; }

        double zeta() const { return _M_zeta; }

        double zeta2theta() const { return _M_zeta2theta; }

        friend bool operator==(const param_type& __p1, const param_type& __p2)
        {
            return __p1._M_a == __p2._M_a
                && __p1._M_b == __p2._M_b
                && __p1._M_theta == __p2._M_theta
                && __p1._M_zeta == __p2._M_zeta
                && __p1._M_zeta2theta == __p2._M_zeta2theta;
        }

    private:
        _IntType _M_a;
        _IntType _M_b;
        double _M_theta;
        double _M_zeta;
        double _M_zeta2theta;

        /**
         * @brief Calculates zeta.
         *
         * @param __n [IN]  The size of the domain.
         * @param __theta [IN]  The skew factor of the distribution.
         */
        double zeta(unsigned long __n, double __theta)
        {
            double ans = 0.0;
            for (unsigned long i=1; i<=__n; ++i)
                ans += std::pow(1.0 / i, __theta);
            return ans;
        }
    };

public:
    /**
     * @brief Constructs a zipfian_int_distribution object.
     *
     * @param __a [IN]  The lower bound of the distribution.
     * @param __b [IN]  The upper bound of the distribution.
     * @param __theta [IN]  The skew factor of the distribution.
     */
    explicit zipfian_int_distribution(_IntType __a = _IntType(0), _IntType __b = _IntType(1), double __theta = 0.99)
        : _M_param(__a, __b, __theta)
    {
    }

    explicit zipfian_int_distribution(const param_type& __p) : _M_param(__p)
    {
    }

    /**
     * @brief Resets the distribution state.
     *
     * Does nothing for the zipfian int distribution.
     */
    void reset() {}

    result_type a() const { return _M_param.a(); }

    result_type b() const { return _M_param.b(); }

    double theta() const { return _M_param.theta(); }

    /**
     * @brief Returns the parameter set of the distribution.
     */
    param_type param() const { return _M_param; }

    /**
     * @brief Sets the parameter set of the distribution.
     * @param __param The new parameter set of the distribution.
     */
    void param(const param_type& __param) { _M_param = __param; }

    /**
     * @brief Returns the inclusive lower bound of the distribution range.
     */
    result_type min() const { return this->a(); }

    /**
     * @brief Returns the inclusive upper bound of the distribution range.
     */
    result_type max() const { return this->b(); }

    /**
     * @brief Generating functions.
     */
    template <typename _UniformRandomNumberGenerator>
    result_type operator()(_UniformRandomNumberGenerator& __urng)
    {
        return this->operator()(__urng, _M_param);
    }

    template <typename _UniformRandomNumberGenerator>
    result_type operator()(_UniformRandomNumberGenerator& __urng, const param_type& __p)
    {
        double alpha = 1 / (1 - __p.theta());
        double eta = (1 - std::pow(2.0 / (__p.b() - __p.a() + 1), 1 - __p.theta())) / (1 - __p.zeta2theta() / __p.zeta());

        double u = std::generate_canonical<double, std::numeric_limits<double>::digits, _UniformRandomNumberGenerator>(__urng);

        double uz = u * __p.zeta();
        if (uz < 1.0)
            return __p.a();
        if (uz < 1.0 + std::pow(0.5, __p.theta()))
            return __p.a() + 1;

        return __p.a() + ((__p.b() - __p.a() + 1) * std::pow(eta * u - eta + 1, alpha));
    }

    /**
     * @brief Return true if two zipfian int distributions have
     *        the same parameters.
     */
    friend bool operator==(const zipfian_int_distribution& __d1, const zipfian_int_distribution& __d2)
    {
        return __d1._M_param == __d2._M_param;
    }

private:
    param_type _M_param;
};